import open3d as o3d
import numpy as np
import imageio
import random
import os


# 环形轨迹，OK
def circular_trajectory(center, up_vector, num_frames, radius):
    trajectory = []
    for i in range(num_frames):
        angle_rad = 2 * np.pi * i / num_frames
        cam_position = np.array([
            radius * np.cos(angle_rad),
            radius * np.sin(angle_rad),
            0
        ])
        direction = (center - cam_position) / np.linalg.norm(center - cam_position)
        trajectory.append({
            "position": cam_position,
            "direction": direction,
            "up": up_vector
        })
    return trajectory

# 抛物线轨迹，不行
def parabolic_trajectory(center, up_vector, num_frames, height, width):
    trajectory = []
    for i in range(num_frames):
        t = i / (num_frames - 1)
        x = -width / 2 + t * width
        z = (4 * height / width**2) * (x**2)  # 抛物线公式
        cam_position = np.array([x, 0, z])
        direction = (center - cam_position) / np.linalg.norm(center - cam_position)
        trajectory.append({
            "position": cam_position,
            "direction": direction,
            "up": up_vector
        })
    return trajectory

# 波浪轨迹，同8字型
def wave_trajectory(center, up_vector, num_frames, amplitude, frequency, radius):
    trajectory = []
    for i in range(num_frames):
        angle_rad = 2 * np.pi * i / num_frames
        height = amplitude * np.sin(frequency * angle_rad)  # 波浪高度
        cam_position = np.array([
            radius * np.cos(angle_rad),
            height,
            radius * np.sin(angle_rad)
        ])
        direction = (center - cam_position) / np.linalg.norm(center - cam_position)
        trajectory.append({
            "position": cam_position,
            "direction": direction,
            "up": up_vector
        })
    return trajectory

# 8字型轨迹，ok
def serpentine_trajectory(center, up_vector, num_frames, amplitude, frequency, radius):
    trajectory = []
    for i in range(num_frames):
        t = i / (num_frames - 1)
        angle_rad = 2 * np.pi * t
        offset = amplitude * np.sin(frequency * angle_rad)
        cam_position = np.array([
            radius * np.cos(angle_rad),
            offset,
            radius * np.sin(angle_rad)
        ])
        direction = (center - cam_position) / np.linalg.norm(center - cam_position)
        trajectory.append({
            "position": cam_position,
            "direction": direction,
            "up": up_vector
        })
    return trajectory

# 螺旋向上轨迹,OK,zoom在2中一直为0
def helical_trajectory(center, up_vector, num_frames, radius, height, rotations):
    trajectory = []
    for i in range(num_frames):
        t = i / (num_frames - 1)
        angle_rad = 2 * np.pi * rotations * t
        radius_new = radius * (1+t)  # radius_new实现半径变化
        # XY平面圆周运动 + Z轴线性上升
        cam_position = np.array([
            radius_new * np.cos(angle_rad),  # X,通过 cos 和 sin 实现圆周运动,半径为radius
            radius_new * np.sin(angle_rad),  # Y
            # height * t**2                # Z轴匀加速运动
            height*t
        ])
        # print(cam_position)
        direction = (center - cam_position) / np.linalg.norm(center - cam_position)
        trajectory.append({
            "position": cam_position,
            "direction": direction,
            "up": up_vector  # 此时 up_vector 应为 [0, 0, 1]
        })
    return trajectory

# 径向远离轨迹，ok，看起来禁止不动（原因是zoom）
def radial_outward_trajectory(center, up_vector, num_frames, start_distance, end_distance):
    trajectory = []
    for i in range(num_frames):
        t = i / (num_frames - 1)
        distance = start_distance + t * (end_distance - start_distance)
        cam_position = center + np.array([distance, 0, 0])
        direction = (center - cam_position) / np.linalg.norm(center - cam_position)
        # print(cam_position,direction)
        trajectory.append({
            "position": cam_position,
            "direction": direction,
            "up": up_vector
        })
    return trajectory


def Render_3DVedio_1(model_path, output_file, trajectory_type, num_frames=90, **kwargs):
    """
    渲染 3D 模型为视频，根据选择的相机轨迹类型生成动画。
    该函数渲染的相机轨道到模型的距离固定。

    参数:
        model_path (str): 3D 模型文件路径。
        output_file (str): 输出视频文件路径。
        trajectory_type (int): 选择轨迹类型 (1-5)。
        num_frames (int): 总帧数。
        **kwargs: 轨迹参数，根据不同轨迹类型需要传入的参数。
    """
    # 设置轨迹中心和上方向
    center = np.array([0, 0, 0])
    up_vector = np.array([0, 0, 1])

    # 定义轨迹类型映射
    trajectory_functions = {
        1: circular_trajectory,
        2: parabolic_trajectory,
        3: wave_trajectory,
        4:  serpentine_trajectory,
    }

    # 检查轨迹类型是否有效
    if trajectory_type not in trajectory_functions:
        raise ValueError("无效的轨迹类型，请选择 1 到 6 之间的数字。")

    # 生成相机轨迹
    trajectory_function = trajectory_functions[trajectory_type]
    trajectory = trajectory_function(center, up_vector, num_frames, **kwargs)

    # 加载模型
    mesh = o3d.io.read_triangle_mesh(model_path)
    if not mesh.has_triangles():
        raise ValueError("模型加载失败或无效，请检查文件路径或格式。")
    mesh.compute_vertex_normals()

    # 创建可视化窗口
    vis = o3d.visualization.Visualizer()
    vis.create_window(width=1920, height=1088, visible=False)

    # 设置背景色为黑色
    render_option = vis.get_render_option()
    render_option.background_color = np.asarray([0, 0, 0])

    # 添加几何体
    vis.add_geometry(mesh)

    # 获取视图控制器
    ctr = vis.get_view_control()
    ctr.set_lookat(center)
    ctr.set_up(up_vector)
    ctr.set_zoom(0.5) # 固定缩放值，渲染出来的视频距离不变

    # 渲染视频帧
    frames = []
    for frame in trajectory:
        ctr.set_front(-frame["direction"])
        ctr.set_up(frame["up"])
        ctr.set_lookat(center)

        vis.poll_events()
        vis.update_renderer()

        screen = vis.capture_screen_float_buffer(do_render=True)
        frame_uint8 = (np.asarray(screen) * 255).astype(np.uint8)
        frames.append(frame_uint8)

    # 关闭窗口
    vis.destroy_window()

    # 保存为视频
    imageio.mimsave(output_file, frames, fps=30)
    print(f"3D 视频已保存为 {output_file}")

def Render_3DVedio_2(model_path, output_file, trajectory_type, num_frames=90, **kwargs):
    """
    渲染 3D 模型为视频，根据选择的相机轨迹类型生成动画。
    该函数渲染的相机轨道到模型的距离变化。

    参数:
        model_path (str): 3D 模型文件路径。
        output_file (str): 输出视频文件路径。
        trajectory_type (int): 选择轨迹类型 (6)。
        num_frames (int): 总帧数。
        **kwargs: 轨迹参数，根据不同轨迹类型需要传入的参数。
    """
    # 设置轨迹中心和上方向
    center = np.array([0, 0, 0])
    up_vector = np.array([0, 0, 1])

    # 定义轨迹类型映射
    trajectory_functions = {
        5: helical_trajectory,
        6: radial_outward_trajectory,
    }

    # 检查轨迹类型是否有效
    if trajectory_type not in trajectory_functions:
        raise ValueError("无效的轨迹类型，请选择 1 到 6 之间的数字。")

    # 生成相机轨迹
    trajectory_function = trajectory_functions[trajectory_type]
    trajectory = trajectory_function(center, up_vector, num_frames, **kwargs)

    # 加载模型
    mesh = o3d.io.read_triangle_mesh(model_path)
    if not mesh.has_triangles():
        raise ValueError("模型加载失败或无效，请检查文件路径或格式。")
    mesh.compute_vertex_normals()

    # 创建可视化窗口
    vis = o3d.visualization.Visualizer()
    vis.create_window(width=1920, height=1088,visible=False)

    # 设置背景色为黑色
    render_option = vis.get_render_option()
    render_option.background_color = np.asarray([0, 0, 0])

    # 添加几何体
    vis.add_geometry(mesh)

    # 获取视图控制器
    ctr = vis.get_view_control()
    ctr.set_lookat(center)
    ctr.set_up(up_vector)
    # 初始化缩放比例参考值
    initial_distance = np.linalg.norm(trajectory[0]["position"] - center)
    # print(initial_distance)
    base_zoom = 0.5  # 初始缩放值

    # 渲染视频帧
    frames = []
    for frame in trajectory:
        position = frame["position"]
        distance = np.linalg.norm(position)  # 当前帧位置的模长作为距离
        zoom = base_zoom * (initial_distance / distance)  # 动态调整缩放比例
        # print("zoom:",zoom)

        ctr.set_front(-frame["direction"])
        ctr.set_up(frame["up"])
        ctr.set_lookat(center)
        ctr.set_zoom(zoom)

        vis.poll_events()
        vis.update_renderer()

        screen = vis.capture_screen_float_buffer(do_render=True)
        frame_uint8 = (np.asarray(screen) * 255).astype(np.uint8)
        frames.append(frame_uint8)

    # 关闭窗口
    vis.destroy_window()

    # 保存为视频
    imageio.mimsave(output_file, frames, fps=30)
    print(f"3D 视频已保存为 {output_file}")

# 主函数
def render_video(model_path: str, output_file: str) -> None:
    # 随机选择轨迹类型
    # trajectory_type = int(input("请输入轨迹类型 (1: 环形, 2: 抛物线, 3: 波浪, 4: 蛇形, 5: 螺旋向上, 6: 径向远离): "))
    trajectory_type = random.randint(1, 6)
    print('trajectory_type:',trajectory_type)

    # 设置模型路径和输出文件路径
    # model_path = "output/0/mesh.obj"
    # output_file = "output/0/trajectory_render.mp4"

    # 定义不同轨迹类型的参数
    trajectory_params_1 = {
        1: {"radius": 10},
        2: {"height": 5, "width": 10},
        3: {"amplitude": 2, "frequency": 3, "radius": 10},
        4: {"amplitude": 2, "frequency": 3, "radius": 10}
    }
    trajectory_params_2 = {
        5:{"radius": 20, "height": 20, "rotations": 0.25},
        6: {"start_distance": 30, "end_distance": 10}
    }

    # 渲染视频
    if trajectory_type in trajectory_params_1:
        Render_3DVedio_1(model_path, output_file, trajectory_type, num_frames=90,
                                        **trajectory_params_1[trajectory_type])
    elif trajectory_type in trajectory_params_2:
        Render_3DVedio_2(model_path, output_file, trajectory_type, num_frames=90,
                                        **trajectory_params_2[trajectory_type])
    else:
        print("数字无效，生成的数字不在 1 到 6 范围内。")

# 调用函数
render_video("output/1/mesh.obj", "output/1/trajectory_render.mp4")